from sys import modules
from django.db.models import Count, Sum
from .models import GPUMachine, StorageMachine
from task.models import Task


def assets_overview():
    storage_count = dict(StorageMachine.objects.values_list(
        'area').annotate(Count('area')))
    gpu_count = dict(GPUMachine.objects.values_list(
        'area').annotate(Count('area')))
    capacity_count = dict(StorageMachine.objects.values_list(
        'area').annotate(Sum('capacity')))
    storage_stop_count = dict(StorageMachine.objects.values_list(
        'node_status').annotate(Count('node_status')))
    storage_error_count = dict(StorageMachine.objects.values_list(
        'node_status').annotate(Count('node_status')))
    gpu_stop_count = dict(GPUMachine.objects.values_list(
        'node_status').annotate(Count('node_status')))
    gpu_error_count = dict(GPUMachine.objects.values_list(
        'node_status').annotate(Count('node_status')))

    return {'storage_count': storage_count,
            'gpu_count': gpu_count,
            'capacity_count': capacity_count,
            'storage_stop_count': storage_stop_count,
            'storage_error_count': storage_error_count,
            'gpu_stop_count': gpu_stop_count,
            'gpu_error_count': gpu_error_count,
            }


def index_overview():
    capacity = StorageMachine.objects.aggregate(
        Sum('capacity'))['capacity__sum']/1000
    task = Task.objects.count()
    disk = StorageMachine.objects.aggregate(Sum('disk'))['disk__sum']/1000
    machine = StorageMachine.objects.count()
    return {'capacity': capacity, 'task': task, 'disk': disk, 'machine': machine}
